from optfunc import *
from scipy.optimize import minimize

def penaltyFunction(f, p, x, epsilon):
    '''
    f: object function
    p: penalty function
    theta: a function about (x, sigma)
    '''
    sigma = 1 # penalty parameter
    while sigma * p(x) > epsilon:
        theta = lambda xk : f(xk) + sigma * p(xk)
        x = unconstrained_optimize(theta, x, 0.001 )

        sigma *= 2
        print(x, sigma)

    return x


def innerPenaltyFunction(f, b, x, epsilon):
    '''
    f: object function
    p: penalty function
    theta: a function about (x, sigma)
    '''
    sigma = 1.0 # penalty parameter
    # epsilon = 1.0 / epsilon
    while sigma * b(x) > epsilon:

        theta = lambda xk : f(xk) + sigma * b(xk)
        x = unconstrained_optimize(theta, x, 0.001 )

        sigma *= 0.7
        print(x, sigma)

    return x


def f1(x):
    x1, x2 = x
    return 0.5 * (x1**2 + 1/3 * x2**2)

def p1(x):
    x1, x2 = x
    return (x1 + x2 - 1)**2


def f2(x):
    x1, x2 = x
    return x1**2 + 4 * x2**2 - 2 * x1 - x2

def st2(x):
    return x[0] - x[1] - 1

def p2(x):
    return st2(x)**2 if st2(x) > 0 else 0

def b2(x):
    return -1 / st2(x)

def f3(x):
    return x[0]**2 + 4 * x[1]**2

def st3(x):
    return np.array([
        st2(x),
        1 - x[0] - x[1]
    ])

def p3(x):
    st = st3(x)
    return np.sum(st[st>0] **2)

def b3(x):
    st = st3(x)
    return np.sum(-1 / st)

# def b3(x):
#     st = -st3(x)
#     return - np.sum(np.log(st))

if __name__ == '__main__':
    x0 = np.array([10, 5])
    
    print("f1:\n外罚函数：")
    penaltyFunction(f1, p1, x0, 0.0001)
    cons = ({'type':'eq', 'fun':lambda x: x[0] + x[1] - 1}, )
    res = minimize(f1, x0, method='SLSQP', constraints=cons)
    print("对照结果",res.x)

    print("f2:\n外罚函数：")
    penaltyFunction(f2, p2, x0, 0.0001)
    print("内罚函数：")
    innerPenaltyFunction(f2, b2, np.array([10, 9.9]), 0.001)
    cons = ({'type':'ineq', 'fun':lambda x: -st2(x)}, )
    res = minimize(f2, x0, method='SLSQP', constraints=cons)
    print("对照结果：", res.x)

    print("f3:")
    penaltyFunction(f3, p3, x0, 0.01)
    print("内罚函数：")
    # innerPenaltyFunction(f3, b3, np.array([2, 1.5]), 0.001)
    cons = ({'type':'ineq', 'fun':lambda x: -st2(x)},
        {'type':'ineq', 'fun':lambda x: x[0] + x[1] - 1}
    )
    res = minimize(f3, x0, method='SLSQP', constraints=cons)
    print("对照结果:", res.x)

